# fp tree数据结构
class TreeNode(object):
    def __init__(self, name_value, num_occur, parent_node):
        # 节点名
        self.name = name_value
        # 出现次数
        self.count = num_occur
        # 链接相似的元素
        self.node_link = None
        self.parent = parent_node
        self.children = {}

    def increase(self, num_occur):
        self.count += num_occur

    def display(self, ind=1):
        """
        展现fp tree
        :param ind:
        :return:
        """
        print('{} {}:{}'.format(' ' * ind, self.name, self.count))
        for child in self.children.values():
            child.display(ind + 1)


class BookFPTree(object):
    @staticmethod
    def update_link(node_to_test, target_node):
        """
        更新 链接相似的元素
        :param node_to_test:
        :param target_node:
        :return:
        """
        while node_to_test.node_link:
            node_to_test = node_to_test.node_link

        node_to_test.node_link = target_node

    def update_tree(self, items, tree, freq_item, count):
        """
        按items[0] 及其频次count 更新 fp tree
        :param items:
        :param tree: fp tree
        :param freq_item: 频繁单元素，以及其对应的次数
        :param count: items[0]的频次
        :return:
        """
        # 若存在，则增加对应值
        if items[0] in tree.children:
            tree.children[items[0]].increase(count)
        else:
            tree.children[items[0]] = TreeNode(items[0], count, parent_node=tree)
            if freq_item[items[0]][1] is None:
                # 设置元素间链接关系
                freq_item[items[0]][1] = tree.children[items[0]]
            else:
                # 更新 链接相似的元素
                self.update_link(freq_item[items[0]][1], tree.children[items[0]])

        if len(items) > 1:
            # 后面的元素，逐个update
            self.update_tree(items[1::], tree.children[items[0]], freq_item, count)

    def create_tree(self, data_set, min_support=1):
        """
        构建fp tree
        :param data_set:
        :param min_support:
        :return:
        """
        item_count = dict()
        for data in data_set:
            for item in data:
                item_count[item] = item_count.get(item, 0) + 1

        freq_item_set = dict()
        for k, v in item_count.items():
            if v >= min_support:
                freq_item_set[k] = v

        freq_item = set(freq_item_set.keys())
        if len(freq_item) == 0:
            return None, None

        # 链接相似的元素，None为无
        for k, v in freq_item_set.items():
            freq_item_set[k] = [v, None]

        tree = TreeNode('Null', 1, None)
        for data, count in data_set.items():
            local_id = {}
            for item in data:
                if item in freq_item:
                    local_id[item] = freq_item_set[item][0]
            if len(local_id) > 0:
                # 避免因为sorted时，存在相同count值的多个元素出现随机排序，故先对元素进行排序
                local_id = sorted(local_id.items(), key=lambda d: d[0])
                ordered_items = [v[0] for v in sorted(local_id, key=lambda d: d[1], reverse=True)]
                print('ordered_items: {}'.format(ordered_items))
                # 按频繁项集排序后的元素，逐个update tree
                self.update_tree(ordered_items, tree, freq_item_set, count)
        return tree, freq_item_set

    def ascend_tree(self, leaf_node, prefix_path):
        """
        从子节点往上迭代整棵树，直到找不到其父节点为止
        :param leaf_node:
        :param prefix_path:
        :return:
        """
        if leaf_node.parent:
            prefix_path.append(leaf_node.name)
            self.ascend_tree(leaf_node.parent, prefix_path)

    def find_prefix_path(self, tree_node):
        """
        从当前值，找出直到根的路径
        :param base_pat:
        :param tree_node:
        :return:
        """
        cond_patds = dict()
        while tree_node:
            prefix_path = []
            # 对每个元素进行往上迭代整棵树，并收集所有遇到元素的名称
            self.ascend_tree(leaf_node=tree_node, prefix_path=prefix_path)
            if len(prefix_path) > 1:
                cond_patds[frozenset(prefix_path[1: ])] = tree_node.count
            tree_node = tree_node.node_link
        return cond_patds

    def mine_tree(self, tree, freq_item_dict, min_support, pre_fix, freq_item_list):
        """
        递归查找频繁项集
        :param tree:
        :param freq_item_dict:
        :param min_support:
        :param pre_fix: 该次处理前的频繁项集list
        :param freq_item_list: 频繁项集列表
        :return:
        """
        freq_item_sort_words = sorted(freq_item_dict.items(), key=lambda d: d[0])
        freq_item_sort_num = [v[0] for v in sorted(freq_item_sort_words, key=lambda d: d[1][0])]
        for base_pat in freq_item_sort_num:
            new_freq_set = pre_fix.copy()
            new_freq_set.add(base_pat)
            freq_item_list.append(new_freq_set)
            con_patt_basess = self.find_prefix_path(tree_node=freq_item_dict[base_pat][1])
            my_cond_tree, my_head = self.create_tree(data_set=con_patt_basess, min_support=min_support)

            print('my_cond_tree: {}'.format(my_cond_tree))
            if my_head:
                print('conditional tree for: {}'.format(new_freq_set))
                my_cond_tree.display(1)
                self.mine_tree(tree=my_cond_tree, freq_item_dict=my_head, min_support=min_support,
                               pre_fix=new_freq_set, freq_item_list=freq_item_list)


class MyFPTree(object):
    def __init__(self, data, min_support):
        self.data = data
        self.min_support = min_support
        self.tree = dict()
        self.leaf_arr = []

    def get_item_freq_sort(self):
        """
        step1: 获取频繁项集，并按支持度降序排列
        :return:
        """
        data_items = [i for j in self.data for i in j]
        support_item = dict()
        for item in set(data_items):
            support_item[item] = data_items.count(item) / len(data)

        # 排序存在的问题：两个item的支持度相同
        support_item_list = sorted(support_item.items(), key=lambda d: d[1], reverse=True)
        return [info[0] for info in support_item_list if info[1] >= self.min_support]

    def data_resort(self, freq_item_list):
        """
        step2: 将原始数据data按频繁项集的顺序重排
        :return:
        """
        new_data = []
        for info in self.data:
            mid_data = [item for item in freq_item_list if item in info]
            new_data.append(mid_data)
        return new_data

    def build_tree_add_item(self, tree, item, direct=None, new_direct=None, num=0, itera_num=0):
        """
        fp tree add 单个元素
        :param tree:
        :param item:
        :return:
        """
        if direct:
            if tree.get(direct, None):
                if item == tree[direct]['node']:
                    tree[direct]['count'] += 1
                else:
                    if itera_num + 1 == num:
                        if tree[direct].get(new_direct, None):

                            if item == tree[direct][new_direct]['node']:
                                tree[direct][new_direct]['count'] += 1
                            else:
                                # 找出item的node方向
                                new_direct = int(new_direct) + 1
                                new_direct_flag = True
                                while new_direct_flag:
                                    if tree[direct].get(str(new_direct), None):
                                        if item == tree[direct][str(new_direct)]['node']:
                                            new_direct_flag = False
                                        else:
                                            new_direct += 1
                                    else:
                                        new_direct_flag = False

                                # 对新方向赋值
                                if tree[direct].get(str(new_direct), None) and (item == tree[direct][str(new_direct)]['node']):
                                    tree[direct][str(new_direct)]['count'] += 1
                                else:
                                    tree[direct][str(new_direct)] = {'node': item, 'count': 1}
                        else:
                            tree[direct][new_direct] = {'node': item, 'count': 1}
                    else:
                        itera_num += 1
                        self.build_tree_add_item(tree[direct], item=item, direct=new_direct, new_direct=new_direct, num=num, itera_num=itera_num)
            else:
                tree[direct] = {'node': item, 'count': 1}
        else:
            direct = '1'
            direct_flag = True
            while direct_flag:
                if not tree.get(direct, None):
                    tree[direct] = {'node': item, 'count': 1}
                    direct_flag = False

                elif tree.get(direct, None) and item == tree[direct]['node']:
                    tree[direct]['count'] += 1
                    direct_flag = False
                else:
                    direct = str(int(direct) + 1)

        if not new_direct:
            new_direct = direct
        else:
            new_direct = str(new_direct)

        return direct, new_direct

    def build_tree_line(self, line):
        """
        按行构建树
        :param tree:
        :param line:
        :return:
        """
        direct = None
        new_direct = '1'
        for i, info in enumerate(line):
            direct, new_direct = self.build_tree_add_item(self.tree, info, direct, new_direct, num=i)

    def create_tree(self):
        """
        构建FP tree
        :param data:
        :return:
        """
        # step1: 获取频繁项集，并按支持度降序排列
        freq_item_list = self.get_item_freq_sort()
        print('freq_item_list: {}'.format(freq_item_list))

        # step2: 将原始数据data按频繁项集的顺序重排
        resort_data = self.data_resort(freq_item_list)
        print('resort_data: {}'.format(resort_data))

        # step3: 构建FP tree
        if not isinstance(resort_data, list):
            resort_data = [resort_data]

        for i, line in enumerate(resort_data):
            if i == 7:
                line = ['I2', 'I1', 'I3', 'I5']
            if i == 8:
                break
            self.build_tree_line(line)

    # tree_dict = {'1': {'node': 'I2', 'count': 7, '1': {'node': 'I1', 'count': 2, '1': {'node': 'I5', 'count': 1}, '2': {'node': 'I4', 'count': 1}}, '2': {'node': 'I4', 'count': 1}, '3': {'node': 'I3', 'count': 4, '3': {'node': 'I1', 'count': 2, '3': {'node': 'I5', 'count': 1}}}}, '2': {'node': 'I3', 'count': 2, '1': {'node': 'I1', 'count': 2}}}
    # print(tree_dict)
    # # 找出 所有频繁项集

    def get_leaf(self, tree_dict):
        """
        找出所有的叶子结点
        :param tree_dict:
        :param leaf_arr:
        :return:
        """

        if len(tree_dict) == 2 and tree_dict.get('node'):
            mid_arr = [tree_dict['node'], tree_dict['count']]
            if mid_arr not in self.leaf_arr:
                self.leaf_arr.append(mid_arr)
        else:
            keys = list(tree_dict.keys())
            if 'node' in keys:
                keys = keys[2:]

            for key in keys:
                self.get_leaf(tree_dict[key])

    # 找出叶子结点

    # 1、树转列表
    # 2、生成项集
    # 3、按最小支持度生成频繁项集


def run_simple():
    # 初始化
    tree_node = TreeNode('pyramid', num_occur=9, parent_node=None)
    # 添加节点
    tree_node.children['eye'] = TreeNode(name_value='eye', num_occur=13, parent_node=None)
    tree_node.children['phoenix'] = TreeNode(name_value='phoenix', num_occur=3, parent_node=None)

    tree_node.display()


def load_simple_data():
    return [['r', 'z', 'h', 'j', 'p'],
            ['z', 'y', 'x', 'w', 'v', 'u', 't', 's'],
            ['z'],
            ['r', 'x', 'n', 'o', 's'],
            ['y', 'r', 'x', 'z', 'q', 't', 'p'],
            ['y', 'z', 'x', 'e', 'q', 's', 't', 'm']]


def create_init_set(data_set):
    data = {}
    for info in data_set:
        data[frozenset(info)] = 1
    return data


def run_book_code():
    """
    book code
    :return:
    """
    data = load_simple_data()
    init_set = create_init_set(data)

    fp_tree = BookFPTree()

    my_tree, freq_item_set = fp_tree.create_tree(init_set, min_support=3)
    my_tree.display()

    # path = fp_tree.find_prefix_path(base_pat='s', tree_node=freq_item_set['s'][1])
    path = fp_tree.find_prefix_path(tree_node=freq_item_set['s'][1])
    print(path)

    fp_tree.mine_tree(tree=my_tree, freq_item_dict=freq_item_set, min_support=3, pre_fix=set(), freq_item_list=[])


def run_my_code():
    """
    my code
    :return:
    """
    data = [['I1', 'I2', 'I5'],
            ['I2', 'I4'],
            ['I2', 'I3'],
            ['I1', 'I2', 'I4'],
            ['I1', 'I3'],
            ['I2', 'I3'],
            ['I1', 'I3'],
            ['I1', 'I2', 'I3', 'I5'],
            ['I1', 'I2', 'I3']]

    min_support_ratio = 2 / len(data)

    my_fptree = MyFPTree(data, min_support=min_support_ratio)
    my_fptree.create_tree()
    tree = my_fptree.tree
    print('FPTree: {}'.format(tree))

    # 找出叶子结点
    my_fptree.get_leaf(tree)
    print('leaf_arr: {}'.format(my_fptree.leaf_arr))


if __name__ == '__main__':
    run_book_code()
    # run_my_code()
